# 2023/7/23 12:39
# xiashuobad
# open_image
# PyCharm
"""
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""

import cv2
import numpy as np
import rasterio
from PIL import Image
from affine import Affine
from osgeo import gdalconst, gdal
from rasterio.windows import Window
from typing import Union
from file_utils import MyPath
from transform_types import geo_transform2affine, format_colormap2rasterio, dtype2gdal_type, \
    format_colormap2pil, reshape_as_image, reshape_as_raster, geo_transform2gdal

Image.MAX_IMAGE_PIXELS = None
gdal.PushErrorHandler('CPLQuietErrorHandler')
gdal.UseExceptions()
# 允许shp属性值中文
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
gdal.SetConfigOption("SHAPE_ENCODING", "GBK")


class OpenImage:
    def __init__(self, fp: str, mode: str = 'r', backend='gdal', **kwargs):
        """
        :param fp: 文件路径
        :param mode: 读写模式，读'r',写'w',更新'u'
        :param backend: 基于那个图片处理库，'gdal','pil','rasterio','cv2'
        :param args:
        :param kwargs:
        """
        self._img_array = None
        self._kwargs = kwargs
        self._filepath = fp
        self._backend = backend.lower()
        self._img_data = None
        self._img_array = None
        self._mode = mode.lower()
        self._geo_transform = None
        self._project = None
        self._shape_type = 'raster' if self._backend in ('gdal', 'rasterio') else 'image'
        self._nodata = None
        self._color_map = None
        if self._mode != 'w':
            assert MyPath(self._filepath).is_file(), f'{self._filepath} 文件路径不存在！'
            self._read_data()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._mode == 'u':
            self.save_update()
        elif self._mode == 'w':
            self.write_data()
        self.close()

    def _read_data(self):
        """
        读取img_data和img_array
        :return:
        """
        if self._backend == 'gdal':
            gdal_access = {'r': gdalconst.GA_ReadOnly, 'u': gdalconst.GA_Update}
            self._img_data = gdal.Open(self._filepath, gdal_access[self._mode])  # type:gdal.Dataset
        elif self._backend == 'pil':
            self._img_data = Image.open(self._filepath, **self._kwargs)  # type:Image.Image
        elif self._backend == 'rasterio':
            self._img_data = rasterio.open(self._filepath, self._mode if self._mode == 'r' else 'r+',
                                           **self._kwargs)  # type:rasterio.DatasetReader
        elif self._backend == 'cv2':
            flags = self._kwargs.get('flags', 1)
            self._img_array = cv2.imdecode(np.fromfile(self._filepath, dtype=np.uint8), flags)
        else:
            raise Exception(f'backend: {self._backend} 暂不支持！')

    def close(self):  # 释放内存资源占用
        """
        释放资源
        :return:
        """
        self._img_data = None
        self._img_array = None

    @property
    def img_data(self) -> Union[gdal.Dataset, Image.Image, rasterio.DatasetReader]:
        """
        返回img_data
        :return:
        """
        return self._img_data

    @property
    def img_array(self) -> np.ndarray:
        """
        返回img_array
        :return:
        """
        if self._img_array is not None:
            return self._img_array
        if isinstance(self.img_data, gdal.Dataset):
            self._img_array = self.img_data.ReadAsArray()
        elif isinstance(self.img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            self._img_array = self.img_data.read()
        elif isinstance(self.img_data, Image.Image):
            self._img_array = np.array(self._img_data)
        return self._img_array

    @property
    def shape(self) -> tuple:
        """
        :return: 返回img_array的shape，(h,w,c) 或者(c,h,w)
        """
        if isinstance(self.img_data, gdal.Dataset):
            channel, height, width = self.img_data.RasterCount, self.img_data.RasterYSize, self.img_data.RasterXSize
            return (height, width) if channel == 1 else (channel, height, width)
        elif isinstance(self.img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            return self.img_data.count, self.img_data.height, self.img_data.width
        else:
            return self.img_array.shape if self.img_array is not None else None

    @property
    def size(self) -> tuple:
        """
        :return: 返回（宽，高），注意是宽在前，高在后
        """
        if self.shape is not None:
            if self._shape_type == 'raster':
                return self.shape[-1], self.shape[-2]
            if self._shape_type == 'image':
                return self.shape[1], self.shape[0]

    @property
    def height(self) -> int:
        """
        返回影像高度
        :return:
        """

        return self.size[1] if self.size is not None else None

    @property
    def width(self) -> int:
        """
        返回影像宽度
        :return:
        """

        return self.size[0] if self.size is not None else None

    @property
    def channel(self) -> int:
        """
        返回影像波段数
        :return:
        """
        if self.shape is not None:
            if len(self.shape) == 2:
                return 1
            if self._shape_type == 'raster':
                return self.shape[0]
            if self._shape_type == 'image':
                return self.shape[-1]

    @property
    def geo_transform(self) -> Union[tuple, Affine]:
        """
        返回影像的坐标系数
        :return:
        """
        if self._geo_transform is not None:
            return self._geo_transform
        if isinstance(self.img_data, gdal.Dataset):
            self._geo_transform = self.img_data.GetGeoTransform()
        elif isinstance(self.img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            self._geo_transform = self.img_data.transform
        return self._geo_transform

    @property
    def project(self):
        """
        返回影像的坐标系
        :return:
        """
        if self._project:
            return self._project
        if isinstance(self.img_data, gdal.Dataset):
            self._project = self.img_data.GetProjection()
        elif isinstance(self.img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            if self._img_data.crs is not None:
                self._project = self.img_data.crs.wkt
        return self._project

    @property
    def dtype(self) -> str:
        """
        返回img_array的dtype类型
        :return:
        """
        if self._img_array is not None:
            return self._img_array.dtype
        return 'uint8'

    @property
    def nodata(self):
        """
        返回影像的无效值
        :return:
        """
        if self._nodata is not None:
            return self._nodata
        if isinstance(self.img_data, gdal.Dataset):
            self._nodata = self.img_data.GetRasterBand(1).GetNoDataValue()
        elif isinstance(self.img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            self._nodata = self.img_data.nodata
        return self._nodata

    @property
    def color_map(self):
        """
        返回影像rgb映射表
        :return:
        """
        return self._color_map

    @property
    def file_path(self):
        """
        返回影像文件路径
        :return:
        """
        return self._filepath

    @property
    def bounds(self) -> list:
        """返回影像的经纬度范围:[x_left,y_bottom,x_right,y_top]"""
        bounds = None
        if isinstance(self._img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            bounds = list(self._img_data.bounds)
        elif isinstance(self._img_data, gdal.Dataset):
            x_left = self.geo_transform[0]
            y_top = self.geo_transform[3]
            x_right = x_left + self.width * self.geo_transform[1] + y_top * self.geo_transform[2]
            y_bottom = y_top + self.height * self.geo_transform[5] + x_left * self.geo_transform[4]
            bounds = [x_left, y_bottom, x_right, y_top]
        return bounds

    def set_geo_transform(self, geo_transform):
        """
        设置影像的坐标系数
        :param geo_transform:
        :return:
        """
        self._geo_transform = geo_transform
        return self

    def set_project(self, project):
        """
        设置影像的坐标系
        :param project:
        :return:
        """
        self._project = project
        return self

    def set_img_array(self, img_array: np.ndarray):
        """
        设置img_array
        :param img_array:
        :return:
        """
        self._img_array = reshape_as_raster(img_array) if self._shape_type == 'raster' else reshape_as_image(img_array)
        return self

    def set_dtype(self, dtype):
        if isinstance(self.img_array, np.ndarray):
            self._img_array.dtype = dtype
        return self

    def array_reverse_rgb(self):
        """
        反转img_array的rgb通道顺序，rgb->bgr,或者bgr->rgb
        :return:
        """
        assert self.img_array is not None, 'img_array is None!'
        if self._shape_type == 'raster':
            self._img_array = self._img_array[::-1, ...]
        else:
            self._img_array = self._img_array[..., ::-1]
        return self

    def _write_gdal_data(self):
        """
        内部私有方法，使用gdal写入影像到文件
        :return:
        """
        if self._color_map is not None and len(self.shape) == 2:
            self._color_map = format_colormap2rasterio(self.color_map)
            color_table = gdal.ColorTable()  # type:gdal.ColorTable
            for i, rgb in self._color_map.items():
                color_table.SetColorEntry(i, rgb)
            self._img_data.GetRasterBand(1).SetColorTable(color_table)
        band_list = list(range(1, self.channel + 1))
        self._img_data.WriteArray(self.img_array, band_list=band_list)

        if self.nodata is not None:
            for i in band_list:
                self.img_data.GetRasterBand(i).SetNoDataValue(self.nodata)
        self._img_data.SetGeoTransform(geo_transform2gdal(self.geo_transform))
        self._img_data.SetProjection(self.project if self.project else '')

    def _write_rasterio_data(self):
        """
        私有方法，使用rasterio写入影像到文件
        :return:
        """
        self._img_data.write(self._img_array,
                             1 if len(self._img_array.shape) == 2 else list(range(1, self.channel + 1)))
        if self.geo_transform is not None:
            self._img_data.transform = self.geo_transform
        if self._project is not None:
            self._img_data.crs = self._project
        if self._color_map is not None and self.channel == 1:
            self._img_data.write_colormap(1, format_colormap2rasterio(self.color_map))

    def save_update(self):
        """
        将影像数据的更新写入到文件
        :return:
        """
        assert self._mode == 'u', f'mode: {self._mode}, 无法更新！'
        if isinstance(self._img_data, gdal.Dataset):
            self._write_gdal_data()
        elif isinstance(self._img_data, rasterio.io.DatasetWriter):
            self._write_rasterio_data()

    def write_data(self):
        """
        将影像数据写入到文件
        :return:
        """
        assert self._mode == 'w', f'mode: {self._mode}, 无法写入！'
        assert self._img_array is not None, f'img_array is None！'
        if self._backend == 'rasterio':
            self.set_geo_transform(geo_transform2affine(self.geo_transform))
            meta = dict(driver='GTiff', width=self.width, height=self.height, count=self.channel, dtype=self.dtype,
                        nodata=self.nodata, transform=self.geo_transform, crs=self.project)
            self._img_data = rasterio.open(self._filepath, 'w', **meta,
                                           **self._kwargs)  # type:rasterio.io.DatasetWriter
            self._write_rasterio_data()
        elif self._backend == 'gdal':
            driver = gdal.GetDriverByName('gtiff')  # type:gdal.Driver
            gdal_data_type = dtype2gdal_type(str(self.dtype))
            self._img_data = driver.Create(self._filepath, self.width, self.height, self.channel, gdal_data_type,
                                           **self._kwargs)  # type:gdal.Dataset
            self._write_gdal_data()

        elif self._backend == 'pil':
            self._img_data = Image.fromarray(self.img_array)
            if self._color_map is not None and len(self.shape) == 2:
                self._img_data.putpalette(format_colormap2pil(self.color_map))
            self._img_data.save(self._filepath)
        elif self._backend == 'cv2':
            cv2.imencode(MyPath(self._filepath).suffix, self.img_array)[1].tofile(self._filepath)

    def set_nodata(self, nodata):
        """
        设置影像的无效值
        :param nodata:
        :return:
        """
        self._nodata = nodata
        return self

    def set_color_map(self, color_map):
        """
        设置影像的rgb映射表
        :param color_map:
        :return:
        """
        self._color_map = color_map
        return self

    def read_window(self, x_off, y_off, window_width, window_height) -> np.ndarray:
        """
        窗口读取模式
        :param x_off: 读取窗口左上角x坐标
        :param y_off: 读取窗口左上角y坐标
        :param window_width: 窗口宽度
        :param window_height: 窗口高度
        :return: 读取到的窗口大小的影像array
        """
        if isinstance(self._img_data, gdal.Dataset):
            return self._img_data.ReadAsArray(x_off, y_off, window_width, window_height)
        elif isinstance(self._img_data, (rasterio.DatasetReader, rasterio.io.DatasetWriter)):
            return self._img_data.read(window=Window(x_off, y_off, window_width, window_height))
        elif self.img_array is not None:  # cv2或者pil
            return self.img_array[x_off:x_off + window_width, y_off:y_off + window_height, ...]

    def show_image(self):
        """
        简单显示一下图片
        """
        assert self.img_array is not None, 'img_array is None!'
        image_show = Image.fromarray(reshape_as_image(self.img_array))
        if self.color_map is not None and (len(self.img_array)) == 2:
            image_show.putpalette(format_colormap2pil(self.color_map))
        image_show.show()


def open_image(fp: str, mode: str = 'r', backend='gdal', **kwargs) -> OpenImage:
    return OpenImage(fp, mode, backend, **kwargs)
